{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG & Embeddings & Vector Store\n",
    "\n",
    "## 💡 这节课会带给你\n",
    "\n",
    "1. 如何用你的垂域数据补充 LLM 的能力\n",
    "2. 如何构建你的垂域（向量）知识库\n",
    "3. 搭建一套完整 RAG 系统需要哪些模块\n",
    "4. 搭建 RAG 系统时更多的有用技巧"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 一、什么是检索增强的生成模型（RAG）\n",
    "\n",
    "### 1.1 大模型目前固有的局限性\n",
    "\n",
    "1. LLM的知识不是实时的\n",
    "2. LLM 可能不知道你私有的领域/业务知识"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 检索增强生成\n",
    "\n",
    "RAG（Retrieval Augmented Generation）顾名思义，通过**检索**的方法来增强**生成模型**的能力。\n",
    "\n",
    "<video src=\"./assets/RAG.mp4\" controls=\"controls\" width=1024px style=\"margin-left: 0px\"></video>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "<b>类比：</b>你可以把这个过程想象成开卷考试。让 LLM 先翻书，再回答问题。\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 二、RAG系统的基本搭建流程\n",
    "\n",
    "<img src=\"./assets/rag.png\" style=\"margin-left: 0px\" width=1024px>\n",
    "\n",
    "搭建过程：\n",
    "\n",
    "1. 文档加载，并按一定条件**切割**成片段\n",
    "2. 将切割的文本片段灌入**检索引擎**\n",
    "3. 封装**检索接口**\n",
    "4. 构建**调用流程**：Query -> 检索 -> Prompt -> LLM -> 回复\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 文档的加载与切割"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install --upgrade openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装 pdf 解析库\n",
    "# !pip install pdfminer.six"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pdfminer.high_level import extract_pages\n",
    "from pdfminer.layout import LTTextContainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1):\n",
    "    '''从 PDF 文件中（按指定页码）提取文字'''\n",
    "    paragraphs = []\n",
    "    buffer = ''\n",
    "    full_text = ''\n",
    "    # 提取全部文本\n",
    "    for i, page_layout in enumerate(extract_pages(filename)):\n",
    "        # 如果指定了页码范围，跳过范围外的页\n",
    "        if page_numbers is not None and i not in page_numbers:\n",
    "            continue\n",
    "        for element in page_layout:\n",
    "            if isinstance(element, LTTextContainer):\n",
    "                full_text += element.get_text() + '\\n'\n",
    "    # 按空行分隔，将文本重新组织成段落\n",
    "    lines = full_text.split('\\n')\n",
    "    for text in lines:\n",
    "        if len(text) >= min_line_length:\n",
    "            buffer += (' '+text) if not text.endswith('-') else text.strip('-')\n",
    "        elif buffer:\n",
    "            paragraphs.append(buffer)\n",
    "            buffer = ''\n",
    "    if buffer:\n",
    "        paragraphs.append(buffer)\n",
    "    return paragraphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Llama 2: Open Foundation and Fine-Tuned Chat Models\n",
      "\n",
      " Hugo Touvron∗ Louis Martin† Kevin Stone† Peter Albert Amjad Almahairi Yasmine Babaei Nikolay Bashlykov Soumya Batra Prajjwal Bhargava Shruti Bhosale Dan Bikel Lukas Blecher Cristian Canton Ferrer Moya Chen Guillem Cucurull David Esiobu Jude Fernandes Jeremy Fu Wenyin Fu Brian Fuller Cynthia Gao Vedanuj Goswami Naman Goyal Anthony Hartshorn Saghar Hosseini Rui Hou Hakan Inan Marcin Kardas Viktor Kerkez Madian Khabsa Isabel Kloumann Artem Korenev Punit Singh Koura Marie-Anne Lachaux Thibaut Lavril Jenya Lee Diana Liskovich Yinghai Lu Yuning Mao Xavier Martinet Todor Mihaylov Pushkar Mishra Igor Molybog Yixin Nie Andrew Poulton Jeremy Reizenstein Rashi Rungta Kalyan Saladi Alan Schelten Ruan Silva Eric Michael Smith Ranjan Subramanian Xiaoqing Ellen Tan Binh Tang Ross Taylor Adina Williams Jian Xiang Kuan Puxin Xu Zheng Yan Iliyan Zarov Yuchen Zhang Angela Fan Melanie Kambadur Sharan Narang Aurelien Rodriguez Robert Stojnic Sergey Edunov Thomas Scialom∗\n",
      "\n",
      " GenAI, Meta\n",
      "\n",
      " In this work, we develop and release Llama 2, a collection of pretrained and ﬁne-tuned large language models (LLMs) ranging in scale from 7 billion to 70 billion parameters. Our ﬁne-tuned LLMs, called Llama 2-Chat, are optimized for dialogue use cases. Our models outperform open-source chat models on most benchmarks we tested, and based onour human evaluations for helpfulness and safety, may be a suitable substitute for closed source models. We provide a detailed description of our approach to ﬁne-tuning and safety improvements of Llama 2-Chat in order to enable the community to build on our work and contribute to the responsible development of LLMs.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "paragraphs = extract_text_from_pdf(\"llama2.pdf\", min_line_length=10)\n",
    "\n",
    "for para in paragraphs[:4]:\n",
    "    print(para+\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 LLM接口封装"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -U python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import os\n",
    "\n",
    "# 加载环境变量\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv(), verbose=True)  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY\n",
    "\n",
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_completion(prompt, model=\"gpt-4o\"):\n",
    "    '''封装 openai 接口'''\n",
    "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    response = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "        temperature=0,  # 模型输出的随机性，0 表示随机性最小\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Prompt 模板"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_prompt(prompt_template, **kwargs):\n",
    "    '''将 Prompt 模板赋值'''\n",
    "    inputs = {}\n",
    "    for k, v in kwargs.items():\n",
    "        if isinstance(v, list) and all(isinstance(elem, str) for elem in v):\n",
    "            val = '\\n\\n'.join(v)\n",
    "        else:\n",
    "            val = v\n",
    "        inputs[k] = val\n",
    "    return prompt_template.format(**inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "你是一个问答机器人。\n",
    "你的任务是根据下述给定的已知信息回答用户问题。\n",
    "\n",
    "已知信息:\n",
    "{context} # 检索出来的原始文档\n",
    "\n",
    "用户问：\n",
    "{query} # 用户的提问\n",
    "\n",
    "如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复\"我无法回答您的问题\"。\n",
    "请不要输出已知信息中不包含的信息或答案。\n",
    "请用中文回答用户问题。\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 三、向量检索\n",
    "\n",
    "### 3.1 什么是向量\n",
    "\n",
    "向量是一种有大小和方向的数学对象。它可以表示为从一个点到另一个点的有向线段。例如，二维空间中的向量可以表示为 $(x,y)$，表示从原点 $(0,0)$ 到点 $(x,y)$ 的有向线段。\n",
    "<br />\n",
    "<img src=\"./assets/vector.png\" style=\"margin-left: 0px\" width=800px>\n",
    "<br />\n",
    "以此类推，我可以用一组坐标 $(x_0, x_1, \\ldots, x_{N-1})$ 表示一个 $N$ 维空间中的向量，$N$ 叫向量的维度。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1.1 文本向量（Text Embeddings）\n",
    "\n",
    "1. 将文本转成一组 $N$ 维浮点数，即**文本向量**又叫 Embeddings\n",
    "2. 向量之间可以计算距离，距离远近对应**语义相似度**大小\n",
    "\n",
    "<br />\n",
    "<img src=\"./assets/embeddings.png\" style=\"margin-left: 0px\" width=800px>\n",
    "<br />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1.2 文本向量是怎么得到的（选）\n",
    "\n",
    "1. 构建相关（正例）与不相关（负例）的句子对样本\n",
    "2. 训练双塔式模型，让正例间的距离小，负例间的距离大\n",
    "\n",
    "例如：\n",
    "\n",
    "<img src=\"./assets/sbert.png\" style=\"margin-left: 0px\" width=500px>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "<b>扩展阅读：https://www.sbert.net</b>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 向量间的相似度计算\n",
    "\n",
    "<img src=\"./assets/sim.png\" style=\"margin-left: 0px\" width=500px>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy import dot\n",
    "from numpy.linalg import norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cos_sim(a, b):\n",
    "    '''余弦距离 -- 越大越相似'''\n",
    "    return dot(a, b)/(norm(a)*norm(b))\n",
    "\n",
    "\n",
    "def l2(a, b):\n",
    "    '''欧氏距离 -- 越小越相似'''\n",
    "    x = np.asarray(a)-np.asarray(b)\n",
    "    return norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embeddings(texts, model=\"text-embedding-ada-002\", dimensions=None):\n",
    "    '''封装 OpenAI 的 Embedding 模型接口'''\n",
    "    if model == \"text-embedding-ada-002\":\n",
    "        dimensions = None\n",
    "    if dimensions:\n",
    "        data = client.embeddings.create(\n",
    "            input=texts, model=model, dimensions=dimensions).data\n",
    "    else:\n",
    "        data = client.embeddings.create(input=texts, model=model).data\n",
    "    return [x.embedding for x in data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total dimension: 1536\n",
      "First 10 elements: [-0.007304091472178698, -0.006229960359632969, -0.010646641254425049, 0.0014391535660251975, -0.010704899206757545, 0.029274623841047287, -0.019807705655694008, 0.005487171467393637, -0.016865678131580353, -0.011979292146861553]\n"
     ]
    }
   ],
   "source": [
    "test_query = [\"测试文本\"]\n",
    "vec = get_embeddings(test_query)[0]\n",
    "print(f\"Total dimension: {len(vec)}\")\n",
    "print(f\"First 10 elements: {vec[:10]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query与自己的余弦距离: 1.00\n",
      "Query与Documents的余弦距离:\n",
      "0.8224810779975097\n",
      "0.8299968969406545\n",
      "0.798096878742543\n",
      "0.7669367418371253\n",
      "0.7933908049643592\n",
      "\n",
      "Query与自己的欧氏距离: 0.00\n",
      "Query与Documents的欧氏距离:\n",
      "0.5958505343143035\n",
      "0.5831005284209486\n",
      "0.6354574886516438\n",
      "0.6827345862824111\n",
      "0.6428206511188914\n"
     ]
    }
   ],
   "source": [
    "query = \"国际争端\"\n",
    "\n",
    "# 且能支持跨语言\n",
    "# query = \"global conflicts\"\n",
    "\n",
    "documents = [\n",
    "    \"联合国就苏丹达尔富尔地区大规模暴力事件发出警告\",\n",
    "    \"土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判\",\n",
    "    \"日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤\",\n",
    "    \"国家游泳中心（水立方）：恢复游泳、嬉水乐园等水上项目运营\",\n",
    "    \"我国首次在空间站开展舱外辐射生物学暴露实验\",\n",
    "]\n",
    "\n",
    "query_vec = get_embeddings([query])[0]\n",
    "doc_vecs = get_embeddings(documents)\n",
    "\n",
    "print(\"Query与自己的余弦距离: {:.2f}\".format(cos_sim(query_vec, query_vec)))\n",
    "print(\"Query与Documents的余弦距离:\")\n",
    "for vec in doc_vecs:\n",
    "    print(cos_sim(query_vec, vec))\n",
    "\n",
    "print()\n",
    "\n",
    "print(\"Query与自己的欧氏距离: {:.2f}\".format(l2(query_vec, query_vec)))\n",
    "print(\"Query与Documents的欧氏距离:\")\n",
    "for vec in doc_vecs:\n",
    "    print(l2(query_vec, vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 向量数据库\n",
    "\n",
    "向量数据库，是专门为向量检索设计的中间件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 为了演示方便，我们只取两页（第一章）\n",
    "paragraphs = extract_text_from_pdf(\n",
    "    \"llama2.pdf\",\n",
    "    page_numbers=[2, 3],\n",
    "    min_line_length=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.config import Settings\n",
    "\n",
    "class MyVectorDBConnector:\n",
    "    def __init__(self, collection_name, embedding_fn):\n",
    "        chroma_client = chromadb.Client(Settings(allow_reset=True))\n",
    "\n",
    "        # 为了演示，实际不需要每次 reset()\n",
    "        chroma_client.reset()\n",
    "\n",
    "        # 创建一个 collection\n",
    "        self.collection = chroma_client.get_or_create_collection(name=collection_name)\n",
    "        self.embedding_fn = embedding_fn\n",
    "\n",
    "    def add_documents(self, documents):\n",
    "        '''向 collection 中添加文档与向量'''\n",
    "        self.collection.add(\n",
    "            embeddings=self.embedding_fn(documents),  # 每个文档的向量\n",
    "            documents=documents,  # 文档的原文\n",
    "            ids=[f\"id{i}\" for i in range(len(documents))]  # 每个文档的 id\n",
    "        )\n",
    "\n",
    "    def search(self, query, top_n):\n",
    "        '''检索向量数据库'''\n",
    "        results = self.collection.query(\n",
    "            query_embeddings=self.embedding_fn([query]),\n",
    "            n_results=top_n\n",
    "        )\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1. Llama 2, an updated version of Llama 1, trained on a new mix of publicly available data. We also increased the size of the pretraining corpus by 40%, doubled the context length of the model, and adopted grouped-query attention (Ainslie et al., 2023). We are releasing variants of Llama 2 with 7B, 13B, and 70B parameters. We have also trained 34B variants, which we report on in this paper but are not releasing.§\n",
      "\n",
      " In this work, we develop and release Llama 2, a family of pretrained and ﬁne-tuned LLMs, Llama 2 and Llama 2-Chat, at scales up to 70B parameters. On the series of helpfulness and safety benchmarks we tested, Llama 2-Chat models generally perform better than existing open-source models. They also appear to be on par with some of the closed-source models, at least on the human evaluations we performed (see Figures 1 and 3). We have taken measures to increase the safety of these models, using safety-speciﬁc data annotation and tuning, as well as conducting red-teaming and employing iterative evaluations. Additionally, this paper contributes a thorough description of our ﬁne-tuning methodology and approach to improving LLM safety. We hope that this openness will enable the community to reproduce ﬁne-tuned LLMs and continue to improve the safety of those models, paving the way for more responsible development of LLMs. We also share novel observations we made during the development of Llama 2 and Llama 2-Chat, such as the emergence of tool usage and temporal organization of knowledge.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 创建一个向量数据库对象\n",
    "vector_db = MyVectorDBConnector(\"demo\", get_embeddings)\n",
    "# 向向量数据库中添加文档\n",
    "vector_db.add_documents(paragraphs)\n",
    "\n",
    "user_query = \"Llama 2有多少参数\"\n",
    "# user_query = \"Does Llama 2 have a conversational variant\"\n",
    "results = vector_db.search(user_query, 2)\n",
    "\n",
    "for para in results['documents'][0]:\n",
    "    print(para+\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "<b>澄清几个关键概念：</b><ul>\n",
    "    <li>向量数据库的意义是快速的检索；</li>\n",
    "    <li>向量数据库本身不生成向量，向量是由 Embedding 模型产生的；</li>\n",
    "    <li>向量数据库与传统的关系型数据库是互补的，不是替代关系，在实际应用中根据实际需求经常同时使用。</li>\n",
    "</ul>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.3.1 向量数据库服务\n",
    "\n",
    "Server 端\n",
    "\n",
    "```sh\n",
    "chroma run --path /db_path\n",
    "```\n",
    "\n",
    "Client 端\n",
    "\n",
    "```python\n",
    "import chromadb\n",
    "chroma_client = chromadb.HttpClient(host='localhost', port=8000)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.3.2 主流向量数据库功能对比\n",
    "\n",
    "<img src=\"./assets/vectordb.png\" style=\"margin-left: 0px\" width=600px>\n",
    "\n",
    "- FAISS: Meta 开源的向量检索引擎 https://github.com/facebookresearch/faiss\n",
    "- Pinecone: 商用向量数据库，只有云服务 https://www.pinecone.io/\n",
    "- **Milvus**: 开源向量数据库，同时有云服务 https://milvus.io/\n",
    "- Weaviate: 开源向量数据库，同时有云服务 https://weaviate.io/\n",
    "- **Qdrant**: 开源向量数据库，同时有云服务 https://qdrant.tech/\n",
    "- PGVector: Postgres 的开源向量检索引擎 https://github.com/pgvector/pgvector\n",
    "- RediSearch: Redis 的开源向量检索引擎 https://github.com/RediSearch/RediSearch\n",
    "- ElasticSearch 也支持向量检索 https://www.elastic.co/enterprise-search/vector-search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "<b>扩展阅读：https://guangzhengli.com/blog/zh/vector-database</b>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.4 基于向量检索的RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RAG_Bot:\n",
    "    def __init__(self, vector_db, llm_api, n_results=2):\n",
    "        self.vector_db = vector_db\n",
    "        self.llm_api = llm_api\n",
    "        self.n_results = n_results\n",
    "\n",
    "    def chat(self, user_query):\n",
    "        # 1. 检索\n",
    "        search_results = self.vector_db.search(user_query, self.n_results)\n",
    "\n",
    "        # 2. 构建 Prompt\n",
    "        prompt = build_prompt(\n",
    "            prompt_template, context=search_results['documents'][0], query=user_query)\n",
    "\n",
    "        # 3. 调用 LLM\n",
    "        response = self.llm_api(prompt)\n",
    "        return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Llama 2有7B、13B和70B三个参数版本。\n"
     ]
    }
   ],
   "source": [
    "# 创建一个RAG机器人\n",
    "bot = RAG_Bot(\n",
    "    vector_db,\n",
    "    llm_api=get_completion\n",
    ")\n",
    "\n",
    "user_query = \"llama 2有多少参数?\"\n",
    "\n",
    "response = bot.chat(user_query)\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.5 OpenAI 新发布的两个 Embedding 模型\n",
    "\n",
    "2024 年 1 月 25 日，OpenAI 新发布了两个 Embedding 模型\n",
    "\n",
    "- text-embedding-3-large\n",
    "- text-embedding-3-small\n",
    "\n",
    "其最大特点是，支持自定义的缩短向量维度，从而在几乎不影响最终效果的情况下降低向量检索与相似度计算的复杂度。\n",
    "\n",
    "通俗的说：**越大越准、越小越快。** 官方公布的评测结果:\n",
    "\n",
    "<img src=\"./assets/mteb.png\" style=\"margin-left: 0px\" width=600px>\n",
    "\n",
    "注：[MTEB](https://huggingface.co/blog/mteb) 是一个大规模多任务的 Embedding 模型公开评测集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "向量维度: 128\n",
      "\n",
      "Query与Documents的余弦距离:\n",
      "0.3341465461188055\n",
      "0.354653344959069\n",
      "0.31360020300733393\n",
      "0.2240495299009075\n",
      "0.12830648514792267\n",
      "\n",
      "Query与Documents的欧氏距离:\n",
      "1.1539960610822753\n",
      "1.13608684105105\n",
      "1.1716653604956542\n",
      "1.245753192131409\n",
      "1.3203738349039935\n"
     ]
    }
   ],
   "source": [
    "model = \"text-embedding-3-large\"\n",
    "dimensions = 128\n",
    "\n",
    "# query = \"国际争端\"\n",
    "\n",
    "# 且能支持跨语言\n",
    "query = \"global conflicts\"\n",
    "\n",
    "documents = [\n",
    "    \"联合国就苏丹达尔富尔地区大规模暴力事件发出警告\",\n",
    "    \"土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判\",\n",
    "    \"日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤\",\n",
    "    \"国家游泳中心（水立方）：恢复游泳、嬉水乐园等水上项目运营\",\n",
    "    \"我国首次在空间站开展舱外辐射生物学暴露实验\",\n",
    "]\n",
    "\n",
    "query_vec = get_embeddings([query], model=model, dimensions=dimensions)[0]\n",
    "doc_vecs = get_embeddings(documents, model=model, dimensions=dimensions)\n",
    "\n",
    "print(\"向量维度: {}\".format(len(query_vec)))\n",
    "\n",
    "print()\n",
    "\n",
    "print(\"Query与Documents的余弦距离:\")\n",
    "for vec in doc_vecs:\n",
    "    print(cos_sim(query_vec, vec))\n",
    "\n",
    "print()\n",
    "\n",
    "print(\"Query与Documents的欧氏距离:\")\n",
    "for vec in doc_vecs:\n",
    "    print(l2(query_vec, vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "<b>扩展阅读：这种可变长度的 Embedding 技术背后的原理叫做 <a href=\"https://arxiv.org/abs/2205.13147\">Matryoshka Representation Learning</a> </b>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 四、实战 RAG 系统的进阶知识"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 文本分割的粒度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**缺陷**\n",
    "\n",
    "1. 粒度太大可能导致检索不精准，粒度太小可能导致信息不全面\n",
    "2. 问题的答案可能跨越两个片段"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个向量数据库对象\n",
    "vector_db = MyVectorDBConnector(\"demo_text_split\", get_embeddings)\n",
    "# 向向量数据库中添加文档\n",
    "vector_db.add_documents(paragraphs)\n",
    "\n",
    "# 创建一个RAG机器人\n",
    "bot = RAG_Bot(\n",
    "    vector_db,\n",
    "    llm_api=get_completion\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " In this work, we develop and release Llama 2, a family of pretrained and ﬁne-tuned LLMs, Llama 2 and Llama 2-Chat, at scales up to 70B parameters. On the series of helpfulness and safety benchmarks we tested, Llama 2-Chat models generally perform better than existing open-source models. They also appear to be on par with some of the closed-source models, at least on the human evaluations we performed (see Figures 1 and 3). We have taken measures to increase the safety of these models, using safety-speciﬁc data annotation and tuning, as well as conducting red-teaming and employing iterative evaluations. Additionally, this paper contributes a thorough description of our ﬁne-tuning methodology and approach to improving LLM safety. We hope that this openness will enable the community to reproduce ﬁne-tuned LLMs and continue to improve the safety of those models, paving the way for more responsible development of LLMs. We also share novel observations we made during the development of Llama 2 and Llama 2-Chat, such as the emergence of tool usage and temporal organization of knowledge.\n",
      "\n",
      " 2. Llama 2-Chat, a ﬁne-tuned version of Llama 2 that is optimized for dialogue use cases. We release\n",
      "\n",
      "====回复====\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'我无法回答您的问题。'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# user_query = \"llama 2有商用许可协议吗\"\n",
    "user_query=\"llama 2 chat有多少参数\"\n",
    "search_results = vector_db.search(user_query, 2)\n",
    "\n",
    "for doc in search_results['documents'][0]:\n",
    "    print(doc+\"\\n\")\n",
    "\n",
    "print(\"====回复====\")\n",
    "bot.chat(user_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Figure 1: Helpfulness human evaluation results for Llama 2-Chat compared to other open-source and closed-source models. Human raters compared model generations on ~4k prompts consisting of both single and multi-turn prompts. The 95% conﬁdence intervals for this evaluation are between 1% and 2%. More details in Section 3.4.2. While reviewing these results, it is important to note that human evaluations can be noisy due to limitations of the prompt set, subjectivity of the review guidelines, subjectivity of individual raters, and the inherent diﬃculty of comparing generations.\n",
      "\n",
      " Figure 2: Win-rate % for helpfulness andsafety between commercial-licensed baselines and Llama 2-Chat, according to GPT 4. To complement the human evaluation, we used a more capable model, not subject to our own guidance. Green area indicates our model is better according to GPT-4. To remove ties, we used win/(win + loss). The orders in which the model responses are presented to GPT-4 are randomly swapped to alleviate bias.\n",
      "\n",
      " 1 Introduction\n",
      "\n",
      " Large Language Models (LLMs) have shown great promise as highly capable AI assistants that excel in complex reasoning tasks requiring expert knowledge across a wide range of ﬁelds, including in specialized domains such as programming and creative writing. They enable interaction with humans through intuitive chat interfaces, which has led to rapid and widespread adoption among the general public.\n",
      "\n",
      " The capabilities of LLMs are remarkable considering the seemingly straightforward nature of the training methodology. Auto-regressive transformers are pretrained on an extensive corpus of self-supervised data, followed by alignment with human preferences via techniques such as Reinforcement Learning with Human Feedback (RLHF). Although the training methodology is simple, high computational requirements have limited the development of LLMs to a few players. There have been public releases of pretrained LLMs (such as BLOOM (Scao et al., 2022), LLaMa-1 (Touvron et al., 2023), and Falcon (Penedo et al., 2023)) that match the performance of closed pretrained competitors like GPT-3 (Brown et al., 2020) and Chinchilla (Hoﬀmann et al., 2022), but none of these models are suitable substitutes for closed “product” LLMs, such as ChatGPT, BARD, and Claude. These closed product LLMs are heavily ﬁne-tuned to align with human preferences, which greatly enhances their usability and safety. This step can require signiﬁcant costs in compute and human annotation, and is often not transparent or easily reproducible, limiting progress within the community to advance AI alignment research.\n",
      "\n",
      " In this work, we develop and release Llama 2, a family of pretrained and ﬁne-tuned LLMs, Llama 2 and Llama 2-Chat, at scales up to 70B parameters. On the series of helpfulness and safety benchmarks we tested, Llama 2-Chat models generally perform better than existing open-source models. They also appear to be on par with some of the closed-source models, at least on the human evaluations we performed (see Figures 1 and 3). We have taken measures to increase the safety of these models, using safety-speciﬁc data annotation and tuning, as well as conducting red-teaming and employing iterative evaluations. Additionally, this paper contributes a thorough description of our ﬁne-tuning methodology and approach to improving LLM safety. We hope that this openness will enable the community to reproduce ﬁne-tuned LLMs and continue to improve the safety of those models, paving the way for more responsible development of LLMs. We also share novel observations we made during the development of Llama 2 and Llama 2-Chat, such as the emergence of tool usage and temporal organization of knowledge.\n",
      "\n",
      "Figure 3: Safety human evaluation results for Llama 2-Chat compared to other open-source and closed source models. Human raters judged model generations for safety violations across ~2,000 adversarial prompts consisting of both single and multi-turn prompts. More details can be found in Section 4.4. It is important to caveat these safety results with the inherent bias of LLM evaluations due to limitations of the prompt set, subjectivity of the review guidelines, and subjectivity of individual raters. Additionally, these safety evaluations are performed using content standards that are likely to be biased towards the Llama 2-Chat models.\n",
      "\n",
      " We are releasing the following models to the general public for research and commercial use‡:\n",
      "\n",
      " 1. Llama 2, an updated version of Llama 1, trained on a new mix of publicly available data. We also increased the size of the pretraining corpus by 40%, doubled the context length of the model, and adopted grouped-query attention (Ainslie et al., 2023). We are releasing variants of Llama 2 with 7B, 13B, and 70B parameters. We have also trained 34B variants, which we report on in this paper but are not releasing.§\n",
      "\n",
      " 2. Llama 2-Chat, a ﬁne-tuned version of Llama 2 that is optimized for dialogue use cases. We release\n",
      "\n",
      " variants of this model with 7B, 13B, and 70B parameters as well.\n",
      "\n",
      " We believe that the open release of LLMs, when done safely, will be a net beneﬁt to society. Like all LLMs, Llama 2 is a new technology that carries potential risks with use (Bender et al., 2021b; Weidinger et al., 2021; Solaiman et al., 2023). Testing conducted to date has been in English and has not — and could not — cover all scenarios. Therefore, before deploying any applications of Llama 2-Chat, developers should perform safety testing and tuning tailored to their speciﬁc applications of the model. We provide a responsible use guide¶ and code examples‖ to facilitate the safe deployment of Llama 2 and Llama 2-Chat. More details of our responsible release strategy can be found in Section 5.3.\n",
      "\n",
      " The remainder of this paper describes our pretraining methodology (Section 2), ﬁne-tuning methodology (Section 3), approach to model safety (Section 4), key observations and insights (Section 5), relevant related work (Section 6), and conclusions (Section 7).\n",
      "\n",
      " ‡https://ai.meta.com/resources/models-and-libraries/llama/ §We are delaying the release of the 34B model due to a lack of time to suﬃciently red team. ¶https://ai.meta.com/llama ‖https://github.com/facebookresearch/llama\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for p in paragraphs:\n",
    "    print(p+\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**改进**: 按一定粒度，部分重叠式的切割文本，使上下文更完整"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.tokenize import sent_tokenize\n",
    "import json\n",
    "\n",
    "\n",
    "def split_text(paragraphs, chunk_size=300, overlap_size=100):\n",
    "    '''按指定 chunk_size 和 overlap_size 交叠割文本'''\n",
    "    sentences = [s.strip() for p in paragraphs for s in sent_tokenize(p)]\n",
    "    chunks = []\n",
    "    i = 0\n",
    "    while i < len(sentences):\n",
    "        chunk = sentences[i]\n",
    "        overlap = ''\n",
    "        prev_len = 0\n",
    "        prev = i - 1\n",
    "        # 向前计算重叠部分\n",
    "        while prev >= 0 and len(sentences[prev])+len(overlap) <= overlap_size:\n",
    "            overlap = sentences[prev] + ' ' + overlap\n",
    "            prev -= 1\n",
    "        chunk = overlap+chunk\n",
    "        next = i + 1\n",
    "        # 向后计算当前chunk\n",
    "        while next < len(sentences) and len(sentences[next])+len(chunk) <= chunk_size:\n",
    "            chunk = chunk + ' ' + sentences[next]\n",
    "            next += 1\n",
    "        chunks.append(chunk)\n",
    "        i = next\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "此处 sent_tokenize 为针对英文的实现，针对中文的实现请参考 chinese_utils.py\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunks = split_text(paragraphs, 300, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个向量数据库对象\n",
    "vector_db = MyVectorDBConnector(\"demo_text_split\", get_embeddings)\n",
    "# 向向量数据库中添加文档\n",
    "vector_db.add_documents(chunks)\n",
    "# 创建一个RAG机器人\n",
    "bot = RAG_Bot(\n",
    "    vector_db,\n",
    "    llm_api=get_completion\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2. Llama 2-Chat, a ﬁne-tuned version of Llama 2 that is optimized for dialogue use cases. We release variants of this model with 7B, 13B, and 70B parameters as well. We believe that the open release of LLMs, when done safely, will be a net beneﬁt to society.\n",
      "\n",
      "In this work, we develop and release Llama 2, a family of pretrained and ﬁne-tuned LLMs, Llama 2 and Llama 2-Chat, at scales up to 70B parameters. On the series of helpfulness and safety benchmarks we tested, Llama 2-Chat models generally perform better than existing open-source models.\n",
      "\n",
      "====回复====\n",
      "Llama 2-Chat有7B、13B和70B三种参数规模。\n"
     ]
    }
   ],
   "source": [
    "# user_query = \"llama 2有商用许可协议吗\"\n",
    "user_query=\"llama 2 chat有多少参数\"\n",
    "\n",
    "search_results = vector_db.search(user_query, 2)\n",
    "for doc in search_results['documents'][0]:\n",
    "    print(doc+\"\\n\")\n",
    "\n",
    "response = bot.chat(user_query)\n",
    "print(\"====回复====\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 检索后排序"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**问题**: 有时，最合适的答案不一定排在检索的最前面"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_query = \"how safe is llama 2\"\n",
    "search_results = vector_db.search(user_query, 5)\n",
    "\n",
    "for doc in search_results['documents'][0]:\n",
    "    print(doc+\"\\n\")\n",
    "\n",
    "response = bot.chat(user_query)\n",
    "print(\"====回复====\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**方案**:\n",
    "\n",
    "1. 检索时过招回一部分文本\n",
    "2. 通过一个排序模型对 query 和 document 重新打分排序\n",
    "\n",
    "<img src=\"./assets/sbert-rerank.png\" style=\"margin-left: 0px\" width=600px>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install sentence_transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import CrossEncoder\n",
    "\n",
    "# model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512) # 英文，模型较小\n",
    "model = CrossEncoder('BAAI/bge-reranker-large', max_length=512) # 多语言，国产，模型较大"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_query = \"how safe is llama 2\"\n",
    "# user_query = \"llama 2安全性如何\"\n",
    "scores = model.predict([(user_query, doc)\n",
    "                       for doc in search_results['documents'][0]])\n",
    "# 按得分排序\n",
    "sorted_list = sorted(\n",
    "    zip(scores, search_results['documents'][0]), key=lambda x: x[0], reverse=True)\n",
    "for score, doc in sorted_list:\n",
    "    print(f\"{score}\\t{doc}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 一些 Rerank 的 API 服务 \n",
    "\n",
    "- [Cohere Rerank](https://cohere.com/rerank)：支持多语言\n",
    "- [Jina Rerank](https://jina.ai/reranker/)：目前只支持英文"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 五、PDF 文档中的表格怎么处理"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./assets/table_rag.png\" style=\"margin-left: 0px\" width=\"800px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 将每页 PDF 转成图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install PyMuPDF\n",
    "# !pip install matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import fitz\n",
    "from PIL import Image\n",
    "\n",
    "def pdf2images(pdf_file):\n",
    "    '''将 PDF 每页转成一个 PNG 图像'''\n",
    "    # 保存路径为原 PDF 文件名（不含扩展名）\n",
    "    output_directory_path, _ = os.path.splitext(pdf_file)\n",
    "    \n",
    "    if not os.path.exists(output_directory_path):\n",
    "        os.makedirs(output_directory_path)\n",
    "    \n",
    "    # 加载 PDF 文件\n",
    "    pdf_document = fitz.open(pdf_file)\n",
    "    \n",
    "    # 每页转一张图\n",
    "    for page_number in range(pdf_document.page_count):\n",
    "        # 取一页\n",
    "        page = pdf_document[page_number]\n",
    "    \n",
    "        # 转图像\n",
    "        pix = page.get_pixmap()\n",
    "    \n",
    "        # 从位图创建 PNG 对象\n",
    "        image = Image.frombytes(\"RGB\", [pix.width, pix.height], pix.samples)\n",
    "    \n",
    "        # 保存 PNG 文件\n",
    "        image.save(f\"./{output_directory_path}/page_{page_number + 1}.png\")\n",
    "    \n",
    "    # 关闭 PDF 文件\n",
    "    pdf_document.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def show_images(dir_path):\n",
    "    '''显示目录下的 PNG 图像'''\n",
    "    for file in os.listdir(dir_path):\n",
    "        if file.endswith('.png'):\n",
    "            # 打开图像\n",
    "            img = Image.open(os.path.join(dir_path, file)) \n",
    "\n",
    "            # 显示图像\n",
    "            plt.imshow(img)\n",
    "            plt.axis('off')  # 不显示坐标轴\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf2images(\"llama2_page8.pdf\")\n",
    "show_images(\"llama2_page8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. 识别文档（图片）中的表格"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaxResize(object):\n",
    "    '''缩放图像'''\n",
    "    def __init__(self, max_size=800):\n",
    "        self.max_size = max_size\n",
    "\n",
    "    def __call__(self, image):\n",
    "        width, height = image.size\n",
    "        current_max_size = max(width, height)\n",
    "        scale = self.max_size / current_max_size\n",
    "        resized_image = image.resize(\n",
    "            (int(round(scale * width)), int(round(scale * height)))\n",
    "        )\n",
    "\n",
    "        return resized_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as transforms\n",
    "\n",
    "# 图像预处理\n",
    "detection_transform = transforms.Compose(\n",
    "    [\n",
    "        MaxResize(800),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForObjectDetection\n",
    "\n",
    "# 加载 TableTransformer 模型\n",
    "model = AutoModelForObjectDetection.from_pretrained(\n",
    "    \"microsoft/table-transformer-detection\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 识别后的坐标换算与后处理\n",
    "\n",
    "def box_cxcywh_to_xyxy(x):\n",
    "    '''坐标转换'''\n",
    "    x_c, y_c, w, h = x.unbind(-1)\n",
    "    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]\n",
    "    return torch.stack(b, dim=1)\n",
    "\n",
    "\n",
    "def rescale_bboxes(out_bbox, size):\n",
    "    '''区域缩放'''\n",
    "    width, height = size\n",
    "    boxes = box_cxcywh_to_xyxy(out_bbox)\n",
    "    boxes = boxes * torch.tensor(\n",
    "        [width, height, width, height], dtype=torch.float32\n",
    "    )\n",
    "    return boxes\n",
    "\n",
    "\n",
    "def outputs_to_objects(outputs, img_size, id2label):\n",
    "    '''从模型输出中取定位框坐标'''\n",
    "    m = outputs.logits.softmax(-1).max(-1)\n",
    "    pred_labels = list(m.indices.detach().cpu().numpy())[0]\n",
    "    pred_scores = list(m.values.detach().cpu().numpy())[0]\n",
    "    pred_bboxes = outputs[\"pred_boxes\"].detach().cpu()[0]\n",
    "    pred_bboxes = [\n",
    "        elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)\n",
    "    ]\n",
    "\n",
    "    objects = []\n",
    "    for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):\n",
    "        class_label = id2label[int(label)]\n",
    "        if not class_label == \"no object\":\n",
    "            objects.append(\n",
    "                {\n",
    "                    \"label\": class_label,\n",
    "                    \"score\": float(score),\n",
    "                    \"bbox\": [float(elem) for elem in bbox],\n",
    "                }\n",
    "            )\n",
    "\n",
    "    return objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# 识别表格，并将表格部分单独存为图像文件\n",
    "\n",
    "def detect_and_crop_save_table(file_path):\n",
    "    # 加载图像（PDF页）    \n",
    "    image = Image.open(file_path)\n",
    "\n",
    "    filename, _ = os.path.splitext(os.path.basename(file_path))\n",
    "\n",
    "    # 输出路径\n",
    "    cropped_table_directory = os.path.join(os.path.dirname(file_path), \"table_images\")\n",
    "\n",
    "    if not os.path.exists(cropped_table_directory):\n",
    "        os.makedirs(cropped_table_directory)\n",
    "\n",
    "    # 预处理\n",
    "    pixel_values = detection_transform(image).unsqueeze(0)\n",
    "\n",
    "    # 识别表格\n",
    "    with torch.no_grad():\n",
    "        outputs = model(pixel_values)\n",
    "\n",
    "    # 后处理，得到表格子区域\n",
    "    id2label = model.config.id2label\n",
    "    id2label[len(model.config.id2label)] = \"no object\"\n",
    "    detected_tables = outputs_to_objects(outputs, image.size, id2label)\n",
    "\n",
    "    print(f\"number of tables detected {len(detected_tables)}\")\n",
    "\n",
    "    for idx in range(len(detected_tables)):\n",
    "        # 将识别从的表格区域单独存为图像\n",
    "        cropped_table = image.crop(detected_tables[idx][\"bbox\"])\n",
    "        cropped_table.save(os.path.join(cropped_table_directory,f\"{filename}_{idx}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detect_and_crop_save_table(\"llama2_page8/page_1.png\")\n",
    "show_images(\"llama2_page8/table_images\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. 基于 GPT-4 Vision API 做表格问答"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI()\n",
    "\n",
    "def encode_image(image_path):\n",
    "  with open(image_path, \"rb\") as image_file:\n",
    "    return base64.b64encode(image_file.read()).decode('utf-8')\n",
    "\n",
    "def image_qa(query, image_path):\n",
    "    base64_image = encode_image(image_path)\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        temperature=0,\n",
    "        seed=42,\n",
    "        messages=[{\n",
    "            \"role\": \"user\",\n",
    "              \"content\": [\n",
    "                  {\"type\": \"text\", \"text\": query},\n",
    "                  {\n",
    "                      \"type\": \"image_url\",\n",
    "                      \"image_url\": {\n",
    "                          \"url\": f\"data:image/jpeg;base64,{base64_image}\",\n",
    "                      },\n",
    "                  },\n",
    "              ],\n",
    "        }],\n",
    "    )\n",
    "\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = image_qa(\"哪个模型在AGI Eval数据集上表现最好。得分多少\",\"llama2_page8/table_images/page_1_0.png\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. 用 GPT-4 Vision 生成表格（图像）描述，并向量化用于检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.config import Settings\n",
    "\n",
    "\n",
    "class NewVectorDBConnector:\n",
    "    def __init__(self, collection_name, embedding_fn):\n",
    "        chroma_client = chromadb.Client(Settings(allow_reset=True))\n",
    "\n",
    "        # 为了演示，实际不需要每次 reset()\n",
    "        chroma_client.reset()\n",
    "\n",
    "        # 创建一个 collection\n",
    "        self.collection = chroma_client.get_or_create_collection(\n",
    "            name=collection_name)\n",
    "        self.embedding_fn = embedding_fn\n",
    "\n",
    "    def add_documents(self, documents):\n",
    "        '''向 collection 中添加文档与向量'''\n",
    "        self.collection.add(\n",
    "            embeddings=self.embedding_fn(documents),  # 每个文档的向量\n",
    "            documents=documents,  # 文档的原文\n",
    "            ids=[f\"id{i}\" for i in range(len(documents))]  # 每个文档的 id\n",
    "        )\n",
    "\n",
    "    def add_images(self, image_paths):\n",
    "        '''向 collection 中添加图像'''\n",
    "        documents = [\n",
    "            image_qa(\"请简要描述图片中的信息\",image)\n",
    "            for image in image_paths\n",
    "        ]\n",
    "        self.collection.add(\n",
    "            embeddings=self.embedding_fn(documents),  # 每个文档的向量\n",
    "            documents=documents,  # 文档的原文\n",
    "            ids=[f\"id{i}\" for i in range(len(documents))],  # 每个文档的 id\n",
    "            metadatas=[{\"image\": image} for image in image_paths] # 用 metadata 标记源图像路径\n",
    "        )\n",
    "\n",
    "    def search(self, query, top_n):\n",
    "        '''检索向量数据库'''\n",
    "        results = self.collection.query(\n",
    "            query_embeddings=self.embedding_fn([query]),\n",
    "            n_results=top_n\n",
    "        )\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = []\n",
    "dir_path = \"llama2_page8/table_images\"\n",
    "for file in os.listdir(dir_path):\n",
    "    if file.endswith('.png'):\n",
    "        # 打开图像\n",
    "        images.append(os.path.join(dir_path, file))\n",
    "\n",
    "new_db_connector = NewVectorDBConnector(\"table_demo\",get_embeddings)\n",
    "new_db_connector.add_images(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query  = \"哪个模型在AGI Eval数据集上表现最好。得分多少\"\n",
    "\n",
    "results = new_db_connector.search(query, 1)\n",
    "metadata = results[\"metadatas\"][0]\n",
    "print(\"====检索结果====\")\n",
    "print(metadata)\n",
    "print(\"====回复====\")\n",
    "response = image_qa(query,metadata[0][\"image\"])\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 一些面向 RAG 的文档解析辅助工具\n",
    "\n",
    "- [PyMuPDF](https://pymupdf.readthedocs.io/en/latest/): PDF 文件处理基础库，带有基于规则的表格与图像抽取（不准）\n",
    "- [RAGFlow](https://github.com/infiniflow/ragflow): 一款基于深度文档理解构建的开源 RAG 引擎，支持多种文档格式\n",
    "- [Unstructured.io](https://unstructured.io/): 一个开源+SaaS形式的文档解析库，支持多种文档格式\n",
    "- [LlamaParse](https://docs.llamaindex.ai/en/stable/llama_cloud/llama_parse/)：付费 API 服务，由 LlamaIndex 官方提供，解析不保证100%准确，实测偶有文字丢失或错位发生\n",
    "- [Mathpix](https://mathpix.com/)：付费 API 服务，效果较好，可解析段落结构、表格、公式等，贵！\n",
    "\n",
    "在工程上，PDF 解析本身是个复杂且琐碎的工作。以上工具都不完美，建议在自己实际场景测试后选择使用。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 六、说说 GraphRAG\n",
    "\n",
    "<img src=\"./assets/GraphRAG.png\" style=\"margin-left: 0px\" width=\"800px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. **什么是 GraphRAG**：核心思想是将知识预先处理成知识图谱\n",
    "2. **优点**：适合复杂问题，尤其是以查询为中心的总结，例如：“XXX团队去年有哪些贡献”\n",
    "3. **缺点**：知识图谱的构建、清洗、维护更新等都有可观的成本\n",
    "4. **建议**：\n",
    "   - GraphRAG 不是万能良药\n",
    "   - 领会其核心思想\n",
    "   - 遇到传统 RAG 无论如何优化都不好解决的问题时，酌情使用"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "\n",
    "### RAG 的流程\n",
    "\n",
    "- 离线步骤：\n",
    "  1. 文档加载\n",
    "  2. 文档切分\n",
    "  3. 向量化\n",
    "  4. 灌入向量数据库\n",
    "- 在线步骤：\n",
    "  1. 获得用户问题\n",
    "  2. 用户问题向量化\n",
    "  3. 检索向量数据库\n",
    "  4. 将检索结果和用户问题填入 Prompt 模版\n",
    "  5. 用最终获得的 Prompt 调用 LLM\n",
    "  6. 由 LLM 生成回复\n",
    "\n",
    "### 我用了一个开源的 RAG，不好使怎么办？\n",
    "\n",
    "1. 检查预处理效果：文档加载是否正确，切割的是否合理\n",
    "2. 测试检索效果：问题检索回来的文本片段是否包含答案\n",
    "3. 测试大模型能力：给定问题和包含答案文本片段的前提下，大模型能不能正确回答问题"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 作业\n",
    "\n",
    "做个自己的 ChatPDF。需求：\n",
    "\n",
    "1. 从本地加载 PDF 文件，基于 PDF 的内容对话\n",
    "2. 可以无前端，只要能在命令行运行就行\n",
    "3. 其它随意发挥"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rag-learn",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
